#ifndef _THREAD_POOL_H_
#define _THREAD_POOL_H_

#include <mutex>
#include <functional>
#include <condition_variable>
#include <thread>
#include <memory>
#include <vector>
#include <atomic>
#include <future>
#include <queue>
#include <utility>

// 以 RAII 手法等待线程结束
class join_thread
{
private:
    std::vector<std::thread>& threads_;
public:
    join_thread(std::vector<std::thread>& threads) : threads_(threads) {}
    ~join_thread()
    {
        for (auto& thread : threads_) {
            if (thread.joinable()) {
                thread.join();
            }
        }
    }
};

class thread_pool
{
private:
    std::atomic<bool> stop;  // 停止标志位
    std::queue<std::function<void()>> task_queue;   // 任务队列
    std::mutex task_mutex;  // 任务队列互斥元
    std::condition_variable task_cond;  // 任务队列条件变量
    std::vector<std::thread> threads;   // 容纳线程对象的容器
    join_thread joiner;

private:
    void work_thread()
    {
        while (!stop) {
            std::unique_lock<std::mutex> lock(task_mutex);
            // 等待任务队列非空
            task_cond.wait(lock, [this]{ return this->stop || !task_queue.empty(); });
            if (stop) break;
            std::function<void()> task = task_queue.front();
            task_queue.pop();
            lock.unlock();  // 不用等到任务完成后再释放锁
            task(); 
        }
    }

public:
    // 返回线程池的全局实例指针
    static thread_pool* globalInstance()
    {
        static thread_pool pool;
        return &pool;
    }

private:
    thread_pool() : stop(false), joiner(threads)
    {
        unsigned int thread_count = std::thread::hardware_concurrency();
        if (thread_count == 0) thread_count = 1;
        try
        {
            for (unsigned int i = 0; i < thread_count; i++) {
                threads.emplace_back(&thread_pool::work_thread, this);
            }
        }
        catch(...)  // 如果线程创建失败, 处理异常
        {
            stop = true;
            task_cond.notify_all();
            throw;  // 将异常抛到上层
        }
    }

    ~thread_pool() 
    { 
        stop = true;
        task_cond.notify_all();
    }

    thread_pool(const thread_pool&) = delete;
    
public:
    template<class Function, class ...Args>
    auto submit(Function&& f, Args&& ...args) -> std::future<decltype(f(args...))>
    {
        using result_type = decltype(f(args...));

        // 因为 std::packaged_task 是可移动不可复制的, 因此用 std::shared_ptr 包装一下
        // 把带参数的任务包装成无参数的任务
        auto task = std::make_shared<std::packaged_task<result_type()>>(
            std::bind(std::forward<Function>(f), std::forward<Args>(args)...)
        );

        // 添加任务到队列中
        {
            std::lock_guard<std::mutex> lock(task_mutex);
            task_queue.emplace([task] {
                (*task)();
            });
        }
        task_cond.notify_one(); // 唤醒一个睡眠的线程去执行任务

        std::future<result_type> res = task->get_future();
        return res;
    }
};

#endif